#!/bin/bash
# AWS环境设置脚本 - 配置AWS环境用于Turritopsis部署

set -e

# 配置参数
AWS_REGION="us-west-2"
KEY_NAME="turritopsis-key"
SECURITY_GROUP="turritopsis-sg"
KEY_FILE="$HOME/.ssh/$KEY_NAME.pem"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_debug() {
    echo -e "${BLUE}[DEBUG]${NC} $1"
}

# 检查AWS CLI
check_aws_cli() {
    log_info "Checking AWS CLI installation..."
    
    if ! command -v aws &> /dev/null; then
        log_error "AWS CLI not found. Please install AWS CLI first:"
        echo "  curl 'https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip' -o 'awscliv2.zip'"
        echo "  unzip awscliv2.zip"
        echo "  sudo ./aws/install"
        exit 1
    fi
    
    # 检查AWS配置
    if ! aws sts get-caller-identity &> /dev/null; then
        log_error "AWS credentials not configured. Please run 'aws configure' first."
        exit 1
    fi
    
    log_info "AWS CLI configured successfully"
    aws sts get-caller-identity
}

# 创建SSH密钥对
create_key_pair() {
    log_info "Creating SSH key pair..."
    
    # 检查密钥是否已存在
    if aws ec2 describe-key-pairs --key-names "$KEY_NAME" --region "$AWS_REGION" &> /dev/null; then
        log_warn "Key pair $KEY_NAME already exists"
        
        # 检查本地密钥文件
        if [ ! -f "$KEY_FILE" ]; then
            log_error "Key pair exists in AWS but local file $KEY_FILE not found"
            log_info "Please download the key pair from AWS Console or delete and recreate it"
            exit 1
        fi
    else
        # 创建新密钥对
        log_info "Creating new key pair: $KEY_NAME"
        aws ec2 create-key-pair \
            --key-name "$KEY_NAME" \
            --region "$AWS_REGION" \
            --query 'KeyMaterial' \
            --output text > "$KEY_FILE"
        
        # 设置正确的权限
        chmod 600 "$KEY_FILE"
        
        log_info "Key pair created and saved to $KEY_FILE"
    fi
}

# 创建安全组
create_security_group() {
    log_info "Creating security group..."
    
    # 检查安全组是否已存在
    if aws ec2 describe-security-groups --group-names "$SECURITY_GROUP" --region "$AWS_REGION" &> /dev/null; then
        log_warn "Security group $SECURITY_GROUP already exists"
        return
    fi
    
    # 创建安全组
    GROUP_ID=$(aws ec2 create-security-group \
        --group-name "$SECURITY_GROUP" \
        --description "Security group for Turritopsis nodes" \
        --region "$AWS_REGION" \
        --output text --query 'GroupId')
    
    log_info "Created security group: $GROUP_ID"
    
    # 添加SSH规则
    aws ec2 authorize-security-group-ingress \
        --group-id "$GROUP_ID" \
        --protocol tcp \
        --port 22 \
        --cidr 0.0.0.0/0 \
        --region "$AWS_REGION"
    
    # 添加Turritopsis端口规则
    aws ec2 authorize-security-group-ingress \
        --group-id "$GROUP_ID" \
        --protocol tcp \
        --port 19001-19020 \
        --cidr 0.0.0.0/0 \
        --region "$AWS_REGION"
    
    # 添加ICMP规则（ping）
    aws ec2 authorize-security-group-ingress \
        --group-id "$GROUP_ID" \
        --protocol icmp \
        --port -1 \
        --cidr 0.0.0.0/0 \
        --region "$AWS_REGION"
    
    log_info "Security group rules configured"
}

# 创建IAM角色（可选）
create_iam_role() {
    log_info "Creating IAM role for EC2 instances..."
    
    ROLE_NAME="TurritopsisEC2Role"
    
    # 检查角色是否已存在
    if aws iam get-role --role-name "$ROLE_NAME" &> /dev/null; then
        log_warn "IAM role $ROLE_NAME already exists"
        return
    fi
    
    # 创建信任策略
    cat > trust-policy.json << EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": {
        "Service": "ec2.amazonaws.com"
      },
      "Action": "sts:AssumeRole"
    }
  ]
}
EOF
    
    # 创建角色
    aws iam create-role \
        --role-name "$ROLE_NAME" \
        --assume-role-policy-document file://trust-policy.json
    
    # 附加基本策略
    aws iam attach-role-policy \
        --role-name "$ROLE_NAME" \
        --policy-arn arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
    
    # 创建实例配置文件
    aws iam create-instance-profile \
        --instance-profile-name "$ROLE_NAME"
    
    aws iam add-role-to-instance-profile \
        --instance-profile-name "$ROLE_NAME" \
        --role-name "$ROLE_NAME"
    
    # 清理临时文件
    rm -f trust-policy.json
    
    log_info "IAM role and instance profile created"
}

# 创建VPC和子网（可选）
create_vpc() {
    log_info "Creating VPC for Turritopsis deployment..."
    
    VPC_NAME="turritopsis-vpc"
    
    # 检查VPC是否已存在
    if aws ec2 describe-vpcs --filters "Name=tag:Name,Values=$VPC_NAME" --region "$AWS_REGION" --query 'Vpcs[0].VpcId' --output text | grep -q "vpc-"; then
        log_warn "VPC $VPC_NAME already exists"
        return
    fi
    
    # 创建VPC
    VPC_ID=$(aws ec2 create-vpc \
        --cidr-block 10.0.0.0/16 \
        --region "$AWS_REGION" \
        --output text --query 'Vpc.VpcId')
    
    # 标记VPC
    aws ec2 create-tags \
        --resources "$VPC_ID" \
        --tags Key=Name,Value="$VPC_NAME" \
        --region "$AWS_REGION"
    
    # 启用DNS解析
    aws ec2 modify-vpc-attribute \
        --vpc-id "$VPC_ID" \
        --enable-dns-hostnames \
        --region "$AWS_REGION"
    
    # 创建互联网网关
    IGW_ID=$(aws ec2 create-internet-gateway \
        --region "$AWS_REGION" \
        --output text --query 'InternetGateway.InternetGatewayId')
    
    # 附加互联网网关到VPC
    aws ec2 attach-internet-gateway \
        --internet-gateway-id "$IGW_ID" \
        --vpc-id "$VPC_ID" \
        --region "$AWS_REGION"
    
    # 创建子网
    SUBNET_ID=$(aws ec2 create-subnet \
        --vpc-id "$VPC_ID" \
        --cidr-block 10.0.1.0/24 \
        --availability-zone "${AWS_REGION}a" \
        --region "$AWS_REGION" \
        --output text --query 'Subnet.SubnetId')
    
    # 创建路由表
    ROUTE_TABLE_ID=$(aws ec2 create-route-table \
        --vpc-id "$VPC_ID" \
        --region "$AWS_REGION" \
        --output text --query 'RouteTable.RouteTableId')
    
    # 添加默认路由
    aws ec2 create-route \
        --route-table-id "$ROUTE_TABLE_ID" \
        --destination-cidr-block 0.0.0.0/0 \
        --gateway-id "$IGW_ID" \
        --region "$AWS_REGION"
    
    # 关联子网和路由表
    aws ec2 associate-route-table \
        --subnet-id "$SUBNET_ID" \
        --route-table-id "$ROUTE_TABLE_ID" \
        --region "$AWS_REGION"
    
    log_info "VPC created: $VPC_ID"
    log_info "Subnet created: $SUBNET_ID"
    log_info "Internet Gateway: $IGW_ID"
}

# 验证设置
verify_setup() {
    log_info "Verifying AWS setup..."
    
    # 检查密钥对
    if aws ec2 describe-key-pairs --key-names "$KEY_NAME" --region "$AWS_REGION" &> /dev/null; then
        log_info "✓ Key pair $KEY_NAME exists"
    else
        log_error "✗ Key pair $KEY_NAME not found"
    fi
    
    # 检查安全组
    if aws ec2 describe-security-groups --group-names "$SECURITY_GROUP" --region "$AWS_REGION" &> /dev/null; then
        log_info "✓ Security group $SECURITY_GROUP exists"
    else
        log_error "✗ Security group $SECURITY_GROUP not found"
    fi
    
    # 检查本地密钥文件
    if [ -f "$KEY_FILE" ]; then
        log_info "✓ Local key file $KEY_FILE exists"
    else
        log_error "✗ Local key file $KEY_FILE not found"
    fi
    
    # 检查SSH连接权限
    if [ -f "$KEY_FILE" ]; then
        chmod 600 "$KEY_FILE"
        log_info "✓ SSH key permissions set correctly"
    fi
}

# 生成配置文件
generate_config() {
    log_info "Generating AWS configuration file..."
    
    cat > aws_config.yaml << EOF
# Turritopsis AWS Configuration
# Generated on: $(date)

aws:
  region: $AWS_REGION
  key_name: $KEY_NAME
  security_group: $SECURITY_GROUP
  instance_type: t3.medium
  ami_id: ami-0c02fb55956c7d316  # Amazon Linux 2

deployment:
  project_dir: /home/ec2-user/turritopsis
  nodes_count: 9
  epochs: 30
  r: 5
  n: 9
  f: 2
  curve: secp256k1
  verify_batch: true
  verify_count: 128
  batch_size: 1000

network:
  router_kind: tcp
  acs_port_range: 19001-19010
  router_port_range: 19010-19020

monitoring:
  collect_logs: true
  collect_metrics: true
  collect_proofs: true
  generate_report: true
EOF
    
    log_info "Configuration saved to aws_config.yaml"
}

# 主函数
main() {
    log_info "Setting up AWS environment for Turritopsis..."
    
    check_aws_cli
    create_key_pair
    create_security_group
    create_iam_role
    create_vpc
    verify_setup
    generate_config
    
    log_info "AWS environment setup completed!"
    log_info "Next steps:"
    log_info "1. Review aws_config.yaml"
    log_info "2. Run ./aws_deploy.sh to create instances"
    log_info "3. Run ./aws_start_turritopsis.sh to start the protocol"
}

# 脚本参数处理
case "${1:-}" in
    "check")
        check_aws_cli
        verify_setup
        ;;
    "key")
        create_key_pair
        ;;
    "sg")
        create_security_group
        ;;
    "vpc")
        create_vpc
        ;;
    "config")
        generate_config
        ;;
    "cleanup")
        log_warn "Cleaning up AWS resources..."
        # 这里可以添加清理逻辑
        log_info "Cleanup completed"
        ;;
    *)
        main
        ;;
esac
